import numpy as np
import heapq
import time

# ---
# 1. 数据生成
#
def generate_vectors(num_vectors=1000, dimensions=2):
    """ 
    生成向量数据
    :param num_vectors: 向量数量
    :param dimensions: 向量维度
    :return: 随机生成的向量数据库
    """ 
    np.random.seed(42)
    return np.random.rand(num_vectors, dimensions)

data_vectors = generate_vectors(num_vectors=1000, dimensions=2)

# ---
# 2. HNSW索引构建
# ---
class HNSW:
    """ 
    实现简单的HNSW索引
    """ 
    def __init__(self, vectors, max_neighbors=5, max_layers=3):
        self.vectors = vectors
        self.max_neighbors = max_neighbors
        self.max_layers = max_layers
        self.layers = self._construct_layers()
    
    def _construct_layers(self):
        layers = []
        current_vectors = self.vectors
        for layer_id in range(self.max_layers):
            layer = {i: [] for i in range(len(current_vectors))}
            for i in range(len(current_vectors)):
                distances = []
                for j in range(len(current_vectors)):
                    if i != j:
                        dist = np.linalg.norm(current_vectors[i] - current_vectors[j])
                        distances.append((dist, j))
                nearest_neighbors = heapq.nsmallest(self.max_neighbors, distances)
                layer[i] = [neighbor[1] for neighbor in nearest_neighbors]
            layers.append(layer)
            # 每层点数减半
            if len(current_vectors) > 1:
                current_vectors = current_vectors[:len(current_vectors)//2]
        return layers

# 构建HNSW索引
start_time = time.time()
hnsw_index = HNSW(data_vectors, max_neighbors=5, max_layers=3)
construction_time = time.time() - start_time

# ---
# 3. HNSW查询实现
# ---
def hnsw_search(query_vector, hnsw_index):
    """ 
    HNSW查询
    :param query_vector: 查询向量
    :param hnsw_index: HNSW索引
    :return: 最近邻索引和距离
    """ 
    current_node = 0
    for layer in range(len(hnsw_index.layers)-1, -1, -1):
        layer_graph = hnsw_index.layers[layer]
        layer_vectors = data_vectors[:len(layer_graph)]
        min_distance = np.linalg.norm(query_vector - layer_vectors[current_node])
        nearest_node = current_node
        
        while True:
            updated = False
            for neighbor in layer_graph[current_node]:
                distance = np.linalg.norm(query_vector - layer_vectors[neighbor])
                if distance < min_distance:
                    min_distance = distance
                    nearest_node = neighbor
                    updated = True
            if not updated:
                break
            current_node = nearest_node
    return nearest_node, min_distance

# 查询最近邻
query_vector = np.array([0.5, 0.5])
start_time = time.time()
nearest_node, distance = hnsw_search(query_vector, hnsw_index)
search_time = time.time() - start_time

# ---
# 4. 输出结果
# ---
print(f"HNSW索引构建时间: {construction_time:.6f} 秒")
print(f"HNSW查询时间: {search_time:.6f} 秒")
print(f"查询向量: {query_vector}")
print(f"最近邻节点索引: {nearest_node}, 距离: {distance:.6f}")